import requests


class APIError(BaseException):
    def __init__(self, error, message, url):
        super(APIError, self).__init__(message)
        self.error = error
        self.url = url


class APIClient:
    base_url = "http://api.zoomeye.org"
    host_query_url = "/host/search"
    web_query_url = "/web/search"
    login_url = "/user/login"
    resource_info_url = "/resources-info"
    retry_times = 3

    def __init__(self):
        self.access_token = None

    def login(self, username, password):
        auth = {'username': username, 'password': password}
        auth_result = requests.post(self.base_url + self.login_url, json=auth).json()
        if 'error' in auth_result.keys():
            raise APIError(message=auth_result['message'], url=auth_result['url'], error=auth_result['error'])
        else:
            self.access_token = auth_result['access_token']

    def resources_info(self):
        headers = {"Authorization": "JWT %s" % self.access_token}
        return requests.get(self.base_url + self.resource_info_url, headers=headers).json()

    def query_host(self, query=None, app=None, version=None, device=None, os=None, service=None, ip=None, cidr=None,
                   hostname=None, port=None, city=None, country=None, asn=None, page=1, facts=None):
        if query is None:
            query = ''
        if app:
            query += ' app:"%s"' % app
        if version:
            query += ' ver:"%s"' % version
        if device:
            query += ' device:"%s"' % device
        if os:
            query += ' os:"%s"' % os
        if service:
            query += ' service:"%s"' % service
        if ip:
            query += ' ip:%s' % ip
        if cidr:
            query += ' cidr:%s' % cidr
        if hostname:
            query += ' hostname:"%s"' % hostname
        if port:
            query += ' port:%s' % port
        if city:
            query += ' city:"%s"' % city
        if country:
            query += ' country:"%s"' % country
        if asn:
            query += ' asn:"%s"' % asn
        query = query.strip()
        return self.__query__(url=self.base_url + self.host_query_url, query=query, page=page, facts=facts)

    def query_web(self, query=None, app=None, header=None, keywords=None, desc=None, title=None, ip=None, site=None,
                  city=None, country=None, page=1, facts=None):
        if query is None:
            query = ''
        if app:
            query += ' app:"%s"' % app
        if header:
            query += ' header:"%s"' % header
        if keywords:
            query += ' keywords:"%s"' % keywords
        if desc:
            query += ' desc:"%s"' % desc
        if title:
            query += ' title:"%s"' % title
        if ip:
            query += ' ip:%s' % ip
        if site:
            query += ' site:%s' % site
        if city:
            query += ' city:"%s"' % city
        if country:
            query += ' country:"%s"' % country
        query = query.strip()
        return self.__query__(url=self.base_url + self.web_query_url, query=query, page=page, facts=facts)

    def __query__(self, url, query, page, facts):
        headers = {"Authorization": "JWT %s" % self.access_token}
        response = requests.get(url, params={"query": query, "page": page, "facts": facts}, headers=headers)
        if response.status_code != 200:
            data = response.json()
            raise APIError(message=data['message'], url=data['url'], error=data['error'])
        return response.json()
